#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
股票一进二板筛选策略

本程序实现了一个完整的股票筛选系统，用于筛选可能从一板晋级到二板的股票
"""

import akshare as ak
import tushare as ts
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import time
import json
import warnings
warnings.filterwarnings('ignore')


class StockSelector:
    """
    股票一进二板筛选器
    """
    
    def __init__(self, config_path="config.json"):
        """
        初始化筛选器
        
        Parameters:
        config_path (str): 配置文件路径
        """
        # 读取配置文件
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = json.load(f)
            
        self.token = self.config.get("tushare_token", "")
        if self.token:
            ts.set_token(self.token)
            self.pro = ts.pro_api()
        else:
            self.pro = None
            
    def get_stock_list(self):
        """
        获取A股股票列表
        
        Returns:
        DataFrame: 股票基本信息
        """
        try:
            # 使用akshare获取股票列表
            stock_list = ak.stock_info_a_code_name()
            return stock_list
        except Exception as e:
            print(f"获取股票列表失败: {e}")
            return pd.DataFrame()
    
    def get_daily_data(self, symbol, start_date, end_date):
        """
        获取股票日线数据
        
        Parameters:
        symbol (str): 股票代码
        start_date (str): 开始日期，格式YYYYMMDD
        end_date (str): 结束日期，格式YYYYMMDD
        
        Returns:
        DataFrame: 日线数据
        """
        try:
            # 使用akshare获取历史数据
            df = ak.stock_zh_a_hist(symbol=symbol, period="daily", 
                                   start_date=start_date, end_date=end_date)
            return df
        except Exception as e:
            print(f"获取{symbol}历史数据失败: {e}")
            return pd.DataFrame()
            
    def get_limit_up_data(self, date):
        """
        获取指定日期涨停板数据
        
        Parameters:
        date (str): 日期，格式YYYYMMDD
        
        Returns:
        DataFrame: 涨停板数据
        """
        try:
            # 使用akshare获取涨停股池数据
            df = ak.stock_zt_pool_em(date=date)
            return df
        except Exception as e:
            print(f"获取{date}涨停数据失败: {e}")
            return pd.DataFrame()
            
    def get_stock_info(self, symbol):
        """
        获取股票基本信息
        
        Parameters:
        symbol (str): 股票代码，如"000001"
        
        Returns:
        dict: 股票基本信息
        """
        try:
            # 使用akshare获取个股信息
            info = ak.stock_individual_info_em(symbol=symbol)
            return info
        except Exception as e:
            print(f"获取{symbol}基本信息失败: {e}")
            return pd.DataFrame()
            
    def calculate_technical_indicators(self, df):
        """
        计算技术指标
        
        Parameters:
        df (DataFrame): 股价数据
        
        Returns:
        DataFrame: 包含技术指标的数据
        """
        # 确保列名正确
        if '收盘' in df.columns:
            close_col = '收盘'
        elif 'close' in df.columns:
            close_col = 'close'
        else:
            close_col = df.columns[2]  # 默认第三列是收盘价
            
        if '最高' in df.columns:
            high_col = '最高'
        elif 'high' in df.columns:
            high_col = 'high'
        else:
            high_col = df.columns[3]  # 默认第四列是最高价
            
        if '最低' in df.columns:
            low_col = '最低'
        elif 'low' in df.columns:
            low_col = 'low'
        else:
            low_col = df.columns[4]  # 默认第五列是最低价
            
        if '成交量' in df.columns:
            volume_col = '成交量'
        elif 'volume' in df.columns:
            volume_col = 'volume'
        else:
            volume_col = df.columns[4]  # 默认第六列是成交量
            
        # 计算移动平均线
        df['MA5'] = df[close_col].rolling(window=5).mean()
        df['MA10'] = df[close_col].rolling(window=10).mean()
        df['MA20'] = df[close_col].rolling(window=20).mean()
        df['MA60'] = df[close_col].rolling(window=60).mean()
        
        # 计算MACD
        df['EMA12'] = df[close_col].ewm(span=12).mean()
        df['EMA26'] = df[close_col].ewm(span=26).mean()
        df['DIF'] = df['EMA12'] - df['EMA26']
        df['DEA'] = df['DIF'].ewm(span=9).mean()
        df['MACD'] = (df['DIF'] - df['DEA']) * 2
        
        # 计算KDJ
        low_list = df[low_col].rolling(9, min_periods=9).min()
        high_list = df[high_col].rolling(9, min_periods=9).max()
        rsv = (df[close_col] - low_list) / (high_list - low_list) * 100
        df['K'] = rsv.ewm(com=2).mean()
        df['D'] = df['K'].ewm(com=2).mean()
        df['J'] = 3 * df['K'] - 2 * df['D']
        
        # 计算换手率（如果不存在）
        if '换手率' not in df.columns and 'turnover' not in df.columns:
            df['换手率'] = 0  # 简化处理，实际应根据流通股本计算
            
        return df
    
    def filter_first_limit_up_features(self, df, limit_info):
        """
        筛选首板特征
        
        Parameters:
        df (DataFrame): 股价数据
        limit_info (dict): 涨停相关信息
        
        Returns:
        bool: 是否符合首板特征
        """
        try:
            # 检查数据是否足够
            if len(df) < 60:
                return False
                
            # 获取最近一个交易日的数据
            latest = df.iloc[-1]
            
            # 1. 涨停时间要早（上午10:30前封板）
            # 这里简化处理，实际需要分时数据
            limit_time = limit_info.get('最后封板时间', '150000')
            if limit_time and len(str(limit_time)) >= 6:
                limit_hour = int(str(limit_time)[:2])
                if limit_hour > self.config.get("limit_time_max", 10):
                    return False
            
            # 2. 封单金额要大（建议流通市值的2%以上）
            # 简化处理，实际需要精确的流通市值数据
            limit_amount = limit_info.get('封单金额', 0)
            if limit_amount < self.config.get("limit_amount_min", 1) * 100000000:  # 1亿
                return False
                
            # 3. 成交量适中（换手率3%-15%）
            if '换手率' in latest:
                turnover = latest['换手率']
                if turnover < self.config.get("turn_over_min", 3) or turnover > self.config.get("turn_over_max", 15):
                    return False
            elif 'turnover' in latest:
                turnover = latest['turnover']
                if turnover < self.config.get("turn_over_min", 3) or turnover > self.config.get("turn_over_max", 15):
                    return False
                    
            # 4. 股价位置：处于相对低位
            # 简化处理：当前价格要小于60日均线
            if 'MA60' in latest and latest['MA60'] > 0:
                if latest['收盘'] > latest['MA60']:
                    return False
                    
            return True
        except Exception as e:
            print(f"筛选首板特征时出错: {e}")
            return False
    
    def filter_technical_indicators(self, df):
        """
        技术面筛选
        
        Parameters:
        df (DataFrame): 包含技术指标的股价数据
        
        Returns:
        bool: 是否符合技术面要求
        """
        try:
            # 检查数据是否足够
            if len(df) < 20:
                return False
                
            # 获取最近一个交易日的数据
            latest = df.iloc[-1]
            prev = df.iloc[-2]
            
            # 1. 均线系统：股价突破重要均线
            if 'MA20' in latest and latest['MA20'] > 0:
                # 股价站上20日均线
                if latest['收盘'] < latest['MA20']:
                    return False
                    
            # 2. MACD金叉或即将金叉
            if 'DIF' in latest and 'DEA' in latest:
                # 当前MACD金叉 或 DEA即将上穿DIF
                if latest['DIF'] < latest['DEA'] and prev['DIF'] < prev['DEA']:
                    # 未出现金叉
                    return False
                    
            # 3. KDJ指标处于强势区域
            if 'K' in latest and 'D' in latest:
                # K值和D值都在50以上，且K>D
                if latest['K'] < 50 or latest['D'] < 50 or latest['K'] < latest['D']:
                    return False
                    
            return True
        except Exception as e:
            print(f"技术面筛选时出错: {e}")
            return False
    
    def filter_fundamental(self, stock_info):
        """
        基本面筛选
        
        Parameters:
        stock_info (DataFrame): 股票基本信息
        
        Returns:
        bool: 是否符合基本面要求
        """
        try:
            # 1. 市值适中（20-200亿）
            # 从stock_info中提取流通市值信息
            market_value = None
            for _, row in stock_info.iterrows():
                if '流通市值' in str(row.get('item', '')) or 'market' in str(row.get('item', '')).lower():
                    try:
                        market_value = float(str(row['value']).replace('亿', ''))
                    except:
                        pass
                    break
                    
            if market_value and (market_value < self.config.get("market_value_min", 20) or 
                                market_value > self.config.get("market_value_max", 200)):
                return False
                
            return True
        except Exception as e:
            print(f"基本面筛选时出错: {e}")
            return True  # 出错时暂时放过
    
    def filter_price_range(self, df):
        """
        根据股价范围筛选
        
        Parameters:
        df (DataFrame): 股价数据
        
        Returns:
        bool: 是否符合价格范围要求
        """
        try:
            if len(df) < 1:
                return False
                
            latest = df.iloc[-1]
            
            # 确定收盘价列名
            if '收盘' in df.columns:
                close_col = '收盘'
            elif 'close' in df.columns:
                close_col = 'close'
            else:
                close_col = df.columns[2]  # 默认第三列是收盘价
                
            price = latest[close_col]
            
            # 判断股价是否在设定范围内
            price_min = self.config.get("price_min", 5)
            price_max = self.config.get("price_max", 100)
            
            if price < price_min or price > price_max:
                return False
                
            return True
        except Exception as e:
            print(f"股价范围筛选时出错: {e}")
            return True  # 出错时暂时放过
    
    def select_stocks(self, date=None):
        """
        筛选符合条件的股票
        
        Parameters:
        date (str): 筛选日期，格式YYYYMMDD，默认为最近一个交易日
        
        Returns:
        DataFrame: 符合条件的股票列表及分析数据
        """
        if not date:
            # 获取最近一个交易日
            date = self.config.get("default_date")
            if not date:
                date = datetime.now().strftime('%Y%m%d')
            
        print(f"开始筛选{date}的股票...")
        
        # 1. 获取涨停板数据
        print("获取涨停板数据...")
        limit_up_df = self.get_limit_up_data(date)
        if limit_up_df.empty:
            print("无法获取涨停数据")
            return pd.DataFrame()
            
        # 筛选首板股票（非连板）
        first_limit_up = limit_up_df[limit_up_df['连板数'] == 1]
        if first_limit_up.empty:
            print("没有首板股票")
            return pd.DataFrame()
            
        selected_stocks = []
        
        # 遍历首板股票进行筛选
        for index, row in first_limit_up.iterrows():
            try:
                symbol = row['代码']
                stock_name = row['名称']
                
                print(f"正在分析 {symbol} {stock_name}...")
                
                # 获取股票历史数据
                end_date = datetime.strptime(date, '%Y%m%d')
                start_date = end_date - timedelta(days=120)  # 获取近4个月数据
                start_date_str = start_date.strftime('%Y%m%d')
                end_date_str = end_date.strftime('%Y%m%d')
                
                hist_df = self.get_daily_data(symbol, start_date_str, end_date_str)
                if hist_df.empty:
                    continue
                    
                # 计算技术指标
                hist_df = self.calculate_technical_indicators(hist_df)
                
                # 1. 首板特征筛选
                if not self.filter_first_limit_up_features(hist_df, row):
                    continue
                    
                # 2. 技术面筛选
                if not self.filter_technical_indicators(hist_df):
                    continue
                    
                # 3. 基本面筛选
                stock_info = self.get_stock_info(symbol)
                if not self.filter_fundamental(stock_info):
                    continue
                    
                # 4. 股价范围筛选
                if not self.filter_price_range(hist_df):
                    continue
                    
                # 符合条件，添加到结果列表
                result = row.to_dict()
                result['分析日期'] = date
                # 添加技术指标
                latest = hist_df.iloc[-1]
                result['收盘价'] = latest.get('收盘', 0)
                result['MA20'] = latest.get('MA20', 0)
                result['MACD'] = latest.get('MACD', 0)
                result['K'] = latest.get('K', 0)
                result['D'] = latest.get('D', 0)
                selected_stocks.append(result)
                
                print(f"  {symbol} {stock_name} 符合条件")
                
                # 控制请求频率，避免被限制
                time.sleep(0.1)
                
            except Exception as e:
                print(f"分析{symbol}时出错: {e}")
                continue
                
        if selected_stocks:
            result_df = pd.DataFrame(selected_stocks)
            return result_df
        else:
            return pd.DataFrame()
    
    def generate_report(self, df):
        """
        生成分析报告
        
        Parameters:
        df (DataFrame): 筛选出的股票数据
        
        Returns:
        str: 分析报告
        """
        if df.empty:
            return "没有符合条件的股票。"
            
        report = []
        report.append("=" * 50)
        report.append("股票一进二板选股报告")
        report.append("=" * 50)
        report.append(f"报告日期: {datetime.now().strftime('%Y-%m-%d')}")
        report.append(f"筛选股票数量: {len(df)}")
        report.append("")
        
        # 列出符合条件的股票
        report.append("符合条件的股票:")
        report.append("-" * 50)
        for index, row in df.iterrows():
            report.append(f"股票代码: {row['代码']}")
            report.append(f"股票名称: {row['名称']}")
            report.append(f"涨停时间: {row.get('最后封板时间', 'N/A')}")
            report.append(f"封单金额: {row.get('封单金额', 'N/A')}")
            report.append(f"换手率: {row.get('换手率', 'N/A')}%")
            report.append(f"收盘价: {row.get('收盘价', 'N/A')}")
            report.append(f"20日均线: {row.get('MA20', 'N/A'):.2f}")
            report.append(f"MACD: {row.get('MACD', 'N/A'):.2f}")
            report.append(f"KDJ: K={row.get('K', 'N/A'):.2f}, D={row.get('D', 'N/A'):.2f}")
            report.append("")
            
        # 风险提示
        report.append("风险提示:")
        report.append("-" * 50)
        report.append("1. 股市有风险，投资需谨慎")
        report.append("2. 以上分析仅供参考，不构成投资建议")
        report.append("3. 建议设置止损位，如跌破5日线止损")
        report.append("4. 控制仓位，避免满仓操作")
        report.append("5. 关注次日开盘表现，如大幅低开应谨慎")
        
        return "\n".join(report)


def main():
    """
    主函数
    """
    # 初始化选股器
    selector = StockSelector()
    
    # 筛选股票
    selected_df = selector.select_stocks()
    
    # 生成报告
    report = selector.generate_report(selected_df)
    print(report)
    
    # 保存结果到文件
    if not selected_df.empty:
        filename = f"selected_stocks_{datetime.now().strftime('%Y%m%d')}.csv"
        selected_df.to_csv(filename, index=False, encoding='utf-8-sig')
        print(f"\n结果已保存到 {filename}")


if __name__ == "__main__":
    main()